// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <cstdint>
#include <vector>

#include <pollux/common/base/bit_util.h>

namespace kumo::pollux {
    /// Dynamic size dense bit set that keeps track of maximum set bit.
    class BitSet {
    public:
        /// Constructs a bitSet. 'min' is the lowest possible member of the set.
        /// Values below this are not present and inserting these is a no-op. 'min' is
        /// used when using this as an IN predicate filter.
        explicit BitSet(int64_t min) : min_(min) {
        }

        void insert(int64_t index) {
            int64_t bit = index - min_;
            if (bit < 0) {
                return;
            }
            if (bit < lastSetBit_) {
                bits::setBit(bits_.data(), bit, true);
                return;
            }
            lastSetBit_ = bit;
            if (lastSetBit_ >= bits_.size() * 64) {
                bits_.resize(std::max<int64_t>(bits_.size() * 2, bits::nwords(bit + 1)));
            }
            bits::setBit(bits_.data(), bit, true);
        }

        bool contains(int64_t index) const {
            auto bit = index - min_;
            if (bit >= bits_.size() * 64) {
                // If index was < min_, bit will have wrapped around and will be >
                // size * 64.
                return false;
            }
            return bits::isBitSet(bits_.data(), bit);
        }

        /// Returns the largest element of the set or 'min_ - 1' if empty.
        int64_t max() const {
            return lastSetBit_ + min_;
        }

        const uint64_t *bits() const {
            return bits_.data();
        }

    private:
        const int64_t min_;
        std::vector<uint64_t> bits_;
        int64_t lastSetBit_ = -1;
    };
} // namespace kumo::pollux
